// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <optional>
#include <set>
#include <unordered_map>

#include <pollux/common/utils.h>
#include <pollux/core/plan_node.h>
#include <pollux/expression/function_signature.h>
#include <pollux/vector/vector_fuzzer.h>

namespace kumo::pollux::plan {
    using fuzzer::DataSpec;

    enum ReferenceQueryErrorCode {
        kSuccess,
        kReferenceQueryFail,
        kReferenceQueryUnsupported
    };

    MELON_ALWAYS_INLINE std::string format_as(ReferenceQueryErrorCode errorCode) {
        switch (errorCode) {
            case ReferenceQueryErrorCode::kSuccess:
                return "kSuccess";
            case ReferenceQueryErrorCode::kReferenceQueryFail:
                return "kReferenceQueryFail";
            case ReferenceQueryErrorCode::kReferenceQueryUnsupported:
                return "kReferenceQueryUnsupported";
            default:
                return "Unknown";
        }
    }

    /// Query runner that uses reference database, i.e. DuckDB, Presto, Spark.
    class ReferenceQueryRunner {
    public:
        enum class RunnerType {
            kPrestoQueryRunner,
            kDuckQueryRunner,
            kSparkQueryRunner
        };

        // @param aggregatePool Used to allocate memory needed for vectors produced
        // by 'execute' methods.
        explicit ReferenceQueryRunner(memory::MemoryPool *aggregatePool)
            : aggregatePool_(aggregatePool) {
        }

        virtual ~ReferenceQueryRunner() = default;

        virtual RunnerType runnerType() const = 0;

        // Scalar types supported by the reference database, to be used to restrict
        // candidates when generating random types for fuzzers.
        virtual const std::vector<TypePtr> &supportedScalarTypes() const {
            return defaultScalarTypes();
        }

        virtual const std::unordered_map<std::string, DataSpec> &
        aggregationFunctionDataSpecs() const = 0;

        /// Converts Pollux plan into SQL accepted by the reference database.
        /// @return std::nullopt if the plan uses features not supported by the
        /// reference database.
        virtual std::optional<std::string> toSql(const core::PlanNodePtr &plan) = 0;

        /// Returns whether a constant expression is supported by the reference
        /// database.
        virtual bool isConstantExprSupported(const core::TypedExprPtr & /*expr*/) {
            return true;
        }

        /// Returns whether types contained in a function signature are all
        /// supported by the reference database.
        virtual bool isSupported(const exec::FunctionSignature & /*signature*/) {
            return true;
        }

        /// Executes SQL query returned by the 'toSql' method using 'input' data.
        /// Converts results using 'resultType' schema.
        virtual std::multiset<std::vector<pollux::variant> > execute(
            const std::string & /*sql*/,
            const std::vector<pollux::RowVectorPtr> & /*input*/,
            const pollux::RowTypePtr & /*resultType*/) {
            POLLUX_UNSUPPORTED();
        }

        // Converts 'plan' into an SQL query and executes it. Result is returned as
        // a MaterializedRowMultiset with the ReferenceQueryErrorCode::kSuccess if
        // successful, or an std::nullopt with a ReferenceQueryErrorCode if the
        // query fails.
        virtual std::pair<std::optional<std::multiset<std::vector<pollux::variant> > >, ReferenceQueryErrorCode>
        execute(const core::PlanNodePtr & /*plan*/) {
            POLLUX_UNSUPPORTED();
        }

        /// Similar to 'execute' but returns results in RowVector format.
        /// Caller should ensure 'supportsPolluxVectorResults' returns true.
        virtual std::pair<
            std::optional<std::vector<pollux::RowVectorPtr> >,
            ReferenceQueryErrorCode>
        executeAndReturnVector(const core::PlanNodePtr & /*plan*/) {
            POLLUX_UNSUPPORTED();
        }

        /// Executes SQL query returned by the 'toSql' method using 'probeInput' and
        /// 'buildInput' data for join node.
        /// Converts results using 'resultType' schema.
        virtual std::multiset<std::vector<pollux::variant> > execute(
            const std::string & /*sql*/,
            const std::vector<RowVectorPtr> & /*probeInput*/,
            const std::vector<RowVectorPtr> & /*buildInput*/,
            const RowTypePtr & /*resultType*/) {
            POLLUX_UNSUPPORTED();
        }

        /// Returns true if 'executeVector' can be called to get results as Pollux
        /// Vector.
        virtual bool supportsPolluxVectorResults() const {
            return false;
        }

        /// Similar to 'execute' but returns results in RowVector format.
        /// Caller should ensure 'supportsPolluxVectorResults' returns true.
        virtual std::vector<RowVectorPtr> executeVector(
            const std::string & /*sql*/,
            const std::vector<RowVectorPtr> & /*input*/,
            const RowTypePtr & /*resultType*/) {
            POLLUX_UNSUPPORTED();
        }

        virtual std::vector<pollux::RowVectorPtr> execute(const std::string & /*sql*/) {
            POLLUX_UNSUPPORTED();
        }

        virtual std::vector<pollux::RowVectorPtr> execute(
            const std::string & /*sql*/,
            const std::string & /*sessionProperty*/) {
            POLLUX_UNSUPPORTED();
        }

        bool isSupportedDwrfType(const TypePtr &type);

        /// Returns the name of the values node table in the form t_<id>.
        std::string getTableName(const core::ValuesNode &valuesNode) {
            return fmt::format("t_{}", valuesNode.id());
        }

    protected:
        memory::MemoryPool *aggregatePool() {
            return aggregatePool_;
        }

        // Traverses all nodes in the plan and returns all tables and their names.
        melon::F14FastMap<std::string, std::vector<pollux::RowVectorPtr> >
        getAllTables(const core::PlanNodePtr &plan);

    private:
        memory::MemoryPool *aggregatePool_;
    };
} // namespace kumo::pollux::plan
